#!/usr/bin/env python
'''
Created on 2/5/15

@author: Rich Plevin (rich@plevin.com)
'''

# Copyright (c) 2015, Richard Plevin.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the
#    distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import sys
import errno
import subprocess
import argparse
from common import readConfigFiles, getParam, CONFIG_VAR_NAME, WORKSPACE_VAR_NAME

PROGRAM = os.path.basename(__file__)
VERSION = "0.5"

CONFIG_FILE_DELIM = ':'


def parseArgs():
    parser = argparse.ArgumentParser(
        prog=PROGRAM,
        description='Queue a GCAM job on a Linux cluster or run the job locally (via "-l" flag)')

    parser.add_argument('-C', '--configFile', type=str, default=None,
                        help='''Specify the one or more GCAM configuration filenames, separated commas.
                                If multiple configuration files are given, the are run in succession in
                                the same "job" on the cluster.
                                N.B. This argument is ignored if scenarios are named via the -s flag.''')

    parser.add_argument('-E', '--enviroVars', type=str, default=None,
                        help='''Comma-delimited list of environment variable assignments to pass
                                to qsub, e.g., -E "FOO=1,BAR=2".''')

    parser.add_argument('-l', '--runLocal', action='store_true', default=False,
                        help="Run GCAM locally on current host, not via qsub")

    parser.add_argument('-m', '--minutes', type=float,
                        help='''Set the number of minutes to allocate for each job submitted.
                                Overrides config parameter GCAM.Minutes.''')

    parser.add_argument('-n', '--noRun', action="store_true",
                        help="Show the command to be run, but don't run it")

    parser.add_argument('-r', '--resources', type=str, default='',
                        help='''Specify resources for the qsub command. Can be a comma-delimited list of
                                assignments NAME=value, e.g., -r pvmem=6GB.''')

    parser.add_argument('-p', '--postProcessor', type=str, default='',
                        help='''Specify the path to a script to run after GCAM completes. It should accept three
                                command-line arguments, the first being the path to the workspace in which GCAM
                                was executed, the second being the name of the configuration file used, and the
                                third being the scenario name of interest. Defaults to value of configuration
                                parameter GCAM.PostProcessor.''')

    parser.add_argument('-P', '--noPostProcessor', action='store_true', default=False,
                        help='''Don't run the post-processor script. (Use this to skip post-processing when a script
                                is named in the ~/.gcam.cfg configuration file.)''')

    parser.add_argument('-Q', '--queueName', type=str, default=None,
                        help='''Specify a queue name for qsub. Default is given by config file
                                param GCAM.DefaultQueue.''')

    parser.add_argument('-s', '--scenario', type=str, default='',
                        help='''Specify the scenario(s) to run. Can be a comma-delimited list of scenario names.
                                The scenarios will be run serially in a single batch job, with an allocated
                                time = GCAM.Minutes * {the number of scenarios}.''')

    parser.add_argument('-S', '--scenariosDir', type=str, default='',
                        help='''Specify the directory holding scenarios. Default is the value of config file param GCAM.ScenariosDir.''')

    parser.add_argument('-V', '--version', action='version', version='%(prog)s ' + VERSION)

    parser.add_argument('-w', '--workspace', type=str, default=None,
                        help='''Specify the path to the GCAM workspace to use. If it doesn't exist, the named workspace
                                will be created. If not specified on the command-line, the value of config file parameter
                                GCAM.Workspace is used, i.e., the "standard" workspace.''')

    args = parser.parse_args()
    return args


def makedirs(dirname, mode=0o770):
    try:
        os.makedirs(dirname, mode=mode)
    except OSError, e:
        if e.errno == errno.EEXIST:
            return
        raise       # otherwise, reraise the error


def setupWorkspace(workspace):
    gcamWorkspace = getParam('GCAM.Workspace')

    def workspaceSymlink(src):
        'Create a link in the new workspace to the equivalent file in the main GCAM workspace'
        dstPath = os.path.join(workspace, src)
        if not os.path.lexists(dstPath):
            dirName = os.path.dirname(dstPath)
            makedirs(dirName)
            srcPath = os.path.join(gcamWorkspace, src)
            os.symlink(srcPath, dstPath)

    # Create the workspace if needed
    if not os.path.isdir(workspace):
        print "Creating GCAM workspace '%s'" % workspace

    # Create a local output dir
    outDir = os.path.join(workspace, 'output')
    makedirs(outDir)

    logPath = os.path.join(workspace, 'exe', 'logs')
    makedirs(logPath)

    # Create link in the new workspace "exe" dir to the executable program and other required files/dirs
    exePath = os.path.join('exe', getParam('GCAM.Executable'))      # expressed as relative to the exe dir
    workspaceSymlink(exePath)
    workspaceSymlink('exe/configuration.xml')       # link to default configuration file
    workspaceSymlink('exe/log_conf.xml')            # and log configuration file
    workspaceSymlink('input')
    workspaceSymlink('libs')


def main():
    args = parseArgs()
    readConfigFiles()

    isQsubbed = (CONFIG_VAR_NAME in os.environ)     # see if we're running as the result of a qsub

    if isQsubbed:
        configFiles = os.environ[CONFIG_VAR_NAME].split(CONFIG_FILE_DELIM)
        workspace   = os.environ[WORKSPACE_VAR_NAME]
        qsub = False
    else:
        scenarios  = args.scenario.split(',') if args.scenario else None
        qsub       = not args.runLocal
        queueName  = args.queueName  or getParam('GCAM.DefaultQueue')
        resources  = args.resources  or getParam('GCAM.QsubResources')
        enviroVars = args.enviroVars or getParam('GCAM.QsubEnviroVars')
        workspace  = args.workspace  or getParam('GCAM.Workspace')
        setupWorkspace(workspace)

    # Optional script to run after successful GCAM runs
    postProcessor = not args.noPostProcessor and (args.postProcessor or getParam('GCAM.PostProcessor'))

    exeDir = os.path.abspath(os.path.join(workspace, 'exe'))

    if not isQsubbed:
        if scenarios:
            # Translate scenario names into config file paths, assuming scenario FOO lives in
            # {scenariosDir}/FOO/config-FOO.xml
            scenariosDir = os.path.abspath(args.scenariosDir) if args.scenariosDir else getParam('GCAM.ScenariosDir')
            configFiles  = map(lambda name: os.path.join(scenariosDir, name, "config-%s.xml" % name), scenarios)
        else:
            configFiles = map(os.path.abspath, args.configFile.split(',')) \
                            if args.configFile else [os.path.join(exeDir, 'configuration.xml')]

    if qsub:
        logFile  = os.path.join(exeDir, 'queueGCAM.log')
        minutes  = (args.minutes or float(getParam('GCAM.Minutes'))) * len(configFiles)
        walltime = "%02d:%02d:00" % (minutes / 60, minutes % 60)

        basicArgs = ['qsub',
                     '-q', queueName,
                     '-N', 'queueGCAM',
                     '-l', 'walltime=%s' % walltime,
                     '-m', 'n',      # never send email
                     '-d', exeDir,
                     '-e', logFile,
                     '-j', 'oe',
                     ]

        # Pass critical args to qsub'd version via environment variables
        configs = CONFIG_FILE_DELIM.join(configFiles)
        qsubEnviroVars = "%s='%s',%s='%s'" % (CONFIG_VAR_NAME, configs, WORKSPACE_VAR_NAME, workspace)
        vFlag = enviroVars + ',' + qsubEnviroVars if enviroVars else qsubEnviroVars

        allArgs = basicArgs + ['-v', vFlag]

        if resources:
            allArgs.extend(['-l', resources])

        scriptPath = os.path.abspath(__file__)
        allArgs.append(scriptPath)       # this script is re-invoked "locally" on compute node
        command = ' '.join(allArgs)
        print command

        if not args.noRun:
            exitStatus = subprocess.call(allArgs, shell=False)
            return exitStatus   # of the qsub command

    else:
        # Run locally, which might mean on a desktop machine, interactively on a
        # compute node (via "qsub -I", or in batch mode on a compute node.
        gcamPath = getParam('GCAM.Executable')
        print "cd", exeDir
        os.chdir(exeDir)        # if isQsubbed, this is redundant but harmless

        exitStatus = 0

        for configFile in configFiles:
            gcamArgs = [gcamPath, '-C%s' % configFile]  # N.B. GCAM doesn't allow space between -C and filename
            command  = ' '.join(gcamArgs)
            print command

            if not args.noRun:
                exitStatus = subprocess.call(gcamArgs, shell=False)
                if exitStatus <> 0:
                    print "GCAM failed with command: %s" % command
                    return exitStatus

                if postProcessor:
                    # use a perl one-liner to extract scenario name from config file
                    perlArgs = ['perl', '-ne', 'print "$1" if /<Value name="scenarioName">([^<]+)</', configFile]
                    scenarioName = subprocess.check_output(perlArgs)
                    args = [postProcessor, workspace, configFile, scenarioName]
                    exitStatus = subprocess.call(args, shell=False)
                    if exitStatus <> 0:
                        print "Post-processor '%s' failed for workspace '%s', configuration file '%s, scenario '%s'" % \
                              (postProcessor, workspace, configFile, scenarioName)
                        return exitStatus

        return exitStatus

if __name__ == '__main__':
    sys.exit(main())
